# -*- coding: UTF-8 -*-
import os
import pprint
import sys
import subprocess
#新增的系统变量使用下面的方法获取不到，关机重启即可
#修改的环境变量是临时改变的，当程序停止时修改的环境变量失效（系统变量不会改变）
# 获取 系统环境 PATH 的变量
env = os.environ.get("PATH")
s1=r"E:\Program Files\HALCON-24.11-Progress-Steady\bin\x64-win64"

# 定义环境变量
os.environ["PATH"] =env+";"+ s1







import halcon as ha
from halcon.numpy_interop import himage_as_numpy_array,himage_from_numpy_array
import math
import os
import sys
import cv2
from PIL import Image
import numpy as np
import pprint
import configparser
import threading
import json
from pathlib import Path
from PyPDF2 import PdfReader, PdfWriter
import fitz
from ocr import tran_en,tran_ch



#清空output文件夹中所有的pdf
path=os.path.dirname(__file__)
output=os.path.join(path,"output")
list_dirs=os.walk(output)
for root, dirs, files in list_dirs:
    for f in files:
        # 分离文件名与扩展名，仅显示txt后缀的文件
        if os.path.splitext(f)[1]=='.pdf':
            file_path=os.path.join(root, f)
            os.remove(file_path)


#打开PDF文件，生成一个对象
file_name="2.pdf"
doc = fitz.open(file_name)
page_numbers = len(doc)  # 获取总页数


default_cover = '1'
cover=input("请输入作为封面模版的页码(默认值为首页)：") or default_cover
page = doc[int(cover)-1]
# 调整缩放因子来使pdf导出更高分辨率的图像。
zoom_x = 3
zoom_y = 3
trans = fitz.Matrix(zoom_x, zoom_y)
pix = page.get_pixmap(matrix=trans, alpha=False)
pil_image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
# 将 Pillow 图像转换为 OpenCV 格式的图像
cv_img = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
h_image=himage_from_numpy_array(cv_img)
ImageEmphasize=ha.emphasize (h_image, 30, 30, 2)

#Image = ha.read_image('D:/桌面/视觉检查/基准书2/封面/3_Page1.jpg')
Width, Height = ha.get_image_size(h_image)


# 创建一个 Halcon 窗口
#WindowHandle = ha.open_window(0, 0, 640, 480, 'black', 'local')

WindowHandle = ha.open_window(0, 0, Width[0]/2, Height[0]/2,father_window=0,mode='visible',machine='')


# 显示图像
#ha.set_display(WindowHandle, 'true_color')
ha.disp_obj(h_image, WindowHandle)

# 设置绘制颜色
ha.set_draw(WindowHandle, 'margin')
ha.set_color(WindowHandle, 'red')  # 设置绘制颜色为红色

print("请框选模版图像roi")
# 通过鼠标绘制矩形 ROI
Row11, Column11, Row21, Column21 = ha.draw_rectangle1(WindowHandle)

# 生成矩形 ROI
ModelRegion = ha.gen_rectangle1(Row11, Column11, Row21, Column21)



print("请框选要识别的文字区域")
ha.set_color(WindowHandle, 'green') 
# 通过鼠标绘制矩形 ROI
Text_Row1, Text_Column1, Text_Row2, Text_Column2 = ha.draw_rectangle1(WindowHandle)
# 生成矩形 ROI
Text_Region = ha.gen_rectangle1(Text_Row1, Text_Column1, Text_Row2, Text_Column2)


# 模型模板
TemplateImage = ha.reduce_domain(ImageEmphasize, ModelRegion)
ModelRegionArea, Model_RefRow, Model_RefColumn=ha.area_center (ModelRegion)


#生成杂波区域
Region_Model_shrink,UsedThreshold=ha.binary_threshold (TemplateImage, 'max_separability', 'dark')
RegionTrans=ha.shape_trans (Region_Model_shrink,'rectangle2')
RegionDilation=ha.dilation_circle (RegionTrans, 3)
ROI_cluster=ha.difference(ModelRegion, RegionDilation)
Area, Row, Column=ha.area_center (ROI_cluster )

# 创建并训练形状模型
ModelID = ha.create_generic_shape_model()
ha.set_generic_shape_model_param(ModelID, 'metric', 'use_polarity')
#ha.set_generic_shape_model_param (ModelID, 'pyramid_level_robust_tracking', 'true')
#ha.set_generic_shape_model_param (ModelID, 'num_levels', 2)
#ha.set_generic_shape_model_param (ModelID, 'contrast_high', 30)
#ha.set_generic_shape_model_param (ModelID, 'contrast_low', 30)
ha.set_generic_shape_model_param(ModelID, 'iso_scale_min', 0.8)

'''
if Area[0] !=0:
    ha.set_generic_shape_model_object (ROI_cluster, ModelID, 'clutter_region')
    ha.set_generic_shape_model_param(ModelID, 'use_clutter', "true")
    ha.set_generic_shape_model_param(ModelID, 'clutter_contrast', 'auto')
    ha.set_generic_shape_model_param(ModelID, 'max_clutter', 0.1)
'''    

ha.write_image (TemplateImage, 'png', 0,"model.png" )
ha.train_generic_shape_model(TemplateImage, ModelID)

# 获取模型轮廓
ModelContours = ha.get_shape_model_contours(ModelID, 1)


# 支持显示模型
Row1, Column1, Row2, Column2 = ha.smallest_rectangle1_xld(ModelContours)
RefRow = (max(Row2) - min(Row1)) / 2
RefColumn = (max(Column2) - min(Column1)) / 2
HomMat2D = ha.vector_angle_to_rigid(0, 0, 0, RefRow, RefColumn, 0)
TransContours = ha.affine_trans_contour_xld(ModelContours, HomMat2D)

# 显示模型轮廓
ha.set_color(WindowHandle, 'green')
ha.set_draw(WindowHandle, 'margin')
ha.disp_obj(TransContours, WindowHandle)
#ha.wait_seconds(1)

# 查找模型
MatchResultID, NumMatchResult = ha.find_generic_shape_model(h_image, ModelID)
Row_model = ha.get_generic_shape_model_result(MatchResultID, 0, 'row')
Column_model = ha.get_generic_shape_model_result(MatchResultID, 0, 'column')
HomMat2D_model = ha.get_generic_shape_model_result(MatchResultID, "best", 'hom_mat_2d')
HomMat2DMatchInvert=ha.hom_mat2d_invert (HomMat2D_model)
print("模版位置：",Row_model,Column_model)

# 设置搜索参数
ha.set_generic_shape_model_param(ModelID, 'angle_start', -math.radians(15))
ha.set_generic_shape_model_param(ModelID, 'angle_end', math.radians(15))
ha.set_generic_shape_model_param(ModelID, 'min_score', 0.65)
ha.set_generic_shape_model_param(ModelID, 'max_overlap', 0)
ha.set_generic_shape_model_param(ModelID, 'border_shape_models', 'false')
ha.set_generic_shape_model_param(ModelID, 'greediness', 0)



#%%
#封面的序列号构成的list
title_index=[]

for i in range(page_numbers):
    page = doc[i]
    trans = fitz.Matrix(zoom_x, zoom_y)
    pix = page.get_pixmap(matrix=trans, alpha=False)
    pil_image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
    # 将 Pillow 图像转换为 OpenCV 格式的图像
    cv_img = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
    h_image=himage_from_numpy_array(cv_img)
    
    #限制目标图像的模版匹配空间，加快搜索速度同时提高鲁棒性
    ModelRegion_expand=ha.dilation_rectangle1 (ModelRegion, 500, 500)
    SearchDomain=ha.reduce_domain(h_image,ModelRegion_expand)
    #锐化
    ImageEmphasize=ha.emphasize (SearchDomain, 30, 30, 2)
    

    # 查找模型
    MatchResultID, NumMatchResult = ha.find_generic_shape_model(ImageEmphasize, ModelID)

    if NumMatchResult >0:
        # 显示检测到的匹配
        ha.disp_obj(h_image, WindowHandle)
        MatchContour = ha.get_generic_shape_model_result_object(MatchResultID, "best", 'contours')
        ha.set_color(WindowHandle, 'green')
        ha.disp_obj(MatchContour, WindowHandle)

        # 检索检测到的匹配的参数
        Row = ha.get_generic_shape_model_result(MatchResultID, "best", 'row')
        Column = ha.get_generic_shape_model_result(MatchResultID, "best", 'column')
        Angle = ha.get_generic_shape_model_result(MatchResultID, "best", 'angle')
        ScaleRow = ha.get_generic_shape_model_result(MatchResultID, "best", 'scale_row')
        ScaleColumn = ha.get_generic_shape_model_result(MatchResultID, "best", 'scale_column')
        HomMat2D = ha.get_generic_shape_model_result(MatchResultID, "best", 'hom_mat_2d')
        
        Score = ha.get_generic_shape_model_result(MatchResultID, "best", 'score')
        print(f"page {i+1},模版匹配得分：{Score}")
            
        #HomMat2D = ha.vector_angle_to_rigid(Row_model, Column_model,0,Row, Column, Angle)
        
        AlignmentHomMat2D=ha.hom_mat2d_compose (HomMat2D, HomMat2DMatchInvert)
        TransformedRegion = ha.affine_trans_region(Text_Region, AlignmentHomMat2D, 'nearest_neighbor')
        ha.disp_region(TransformedRegion, WindowHandle)


        reduce_Image = ha.reduce_domain(h_image, TransformedRegion)

        ImagePart=ha.crop_domain(reduce_Image)
        cropImg=himage_as_numpy_array(ImagePart)
        #cv2.imwrite('error.jpg',cropImg)
        result,img=tran_en(cropImg)
        if len(result) !=0:
            #print("识别到文本：",result[0][1])
            #合并多行文本
            extracted_list = [item[1][0] for item in result]
            text=" ".join(extracted_list)
            #替换ocr容易识别错误的字符
            text=text.replace("O","0")
            text=text.replace("I","1")
            title_index.append([i,text])
            
        # 暂停以便查看结果
        ha.wait_seconds(1)

# 关闭窗口
ha.close_window(WindowHandle)

print("找到文件份数:",len(title_index))
pprint.pprint(title_index)



#提取pdf的i~f页生成一个新的pdf
def gen_pdf(pdf,start,end):
    pdf_new = PdfWriter()
    for i in range(start,end):
        pdf_new.add_page(pdf.pages[i])
    return pdf_new


path=os.path.dirname(__file__)
output=os.path.join(path,"output")

pdf = PdfReader(file_name)
for index,j in enumerate(title_index):
    if index<len(title_index)-1:
        start=title_index[index][0]
        end=title_index[index+1][0]
    else:
        start=title_index[index][0]
        end=page_numbers
    part_pdf=gen_pdf(pdf,start,end)
    savepath=os.path.join(output,f"{j[1]}.pdf")
    print("生成pdf文件:",savepath)
    file=open(savepath,'wb')
    part_pdf.write(file)
    file.close()

